import time
import threading
from typing import Optional

class SnowflakeGenerator:
    """
    雪花算法ID生成器
    生成15位左右的唯一ID
    """
    
    def __init__(self, datacenter_id: int = 1, worker_id: int = 1):
        """
        初始化雪花算法生成器
        
        Args:
            datacenter_id: 数据中心ID (0-3)
            worker_id: 工作机器ID (0-3)
        """
        # 时间戳位数 (24位，可以使用约16年)
        self.TIMESTAMP_BITS = 24
        # 数据中心位数 (2位，最多4个数据中心)
        self.DATACENTER_BITS = 2
        # 工作机器位数 (2位，每个数据中心最多4台机器)
        self.WORKER_BITS = 2
        # 序列号位数 (7位，每毫秒最多128个ID)
        self.SEQUENCE_BITS = 7
        
        # 最大值
        self.MAX_DATACENTER_ID = (1 << self.DATACENTER_BITS) - 1
        self.MAX_WORKER_ID = (1 << self.WORKER_BITS) - 1
        self.MAX_SEQUENCE = (1 << self.SEQUENCE_BITS) - 1
        
        # 偏移量
        self.DATACENTER_SHIFT = self.SEQUENCE_BITS + self.WORKER_BITS
        self.TIMESTAMP_SHIFT = self.SEQUENCE_BITS + self.WORKER_BITS + self.DATACENTER_BITS
        
        # 验证参数
        if datacenter_id > self.MAX_DATACENTER_ID or datacenter_id < 0:
            raise ValueError(f"数据中心ID必须在0-{self.MAX_DATACENTER_ID}之间")
        if worker_id > self.MAX_WORKER_ID or worker_id < 0:
            raise ValueError(f"工作机器ID必须在0-{self.MAX_WORKER_ID}之间")
        
        self.datacenter_id = datacenter_id
        self.worker_id = worker_id
        
        # 序列号
        self.sequence = 0
        # 上次生成ID的时间戳
        self.last_timestamp = -1
        
        # 线程锁
        self.lock = threading.Lock()
        
        # 起始时间戳 (2024-01-01 00:00:00 UTC) - 更近的时间戳，减少位数
        self.EPOCH = 1704067200000
    
    def _get_timestamp(self) -> int:
        """获取当前时间戳（毫秒）"""
        return int(time.time() * 1000)
    
    def _wait_for_next_millisecond(self, last_timestamp: int) -> int:
        """等待到下一毫秒"""
        timestamp = self._get_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._get_timestamp()
        return timestamp
    
    def generate_id(self) -> str:
        """
        生成雪花算法ID
        
        Returns:
            str: 15位左右的唯一ID字符串
        """
        with self.lock:
            timestamp = self._get_timestamp()
            
            # 如果当前时间小于上次生成ID的时间，说明系统时钟回退了
            if timestamp < self.last_timestamp:
                raise RuntimeError("系统时钟回退，拒绝生成ID")
            
            # 如果是同一毫秒内生成ID
            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE
                # 如果序列号溢出，等待下一毫秒
                if self.sequence == 0:
                    timestamp = self._wait_for_next_millisecond(self.last_timestamp)
            else:
                # 不同毫秒，序列号重置为0
                self.sequence = 0
            
            self.last_timestamp = timestamp
            
            # 生成ID
            id_value = (
                ((timestamp - self.EPOCH) << self.TIMESTAMP_SHIFT) |
                (self.datacenter_id << self.DATACENTER_SHIFT) |
                (self.worker_id << self.SEQUENCE_BITS) |
                self.sequence
            )
            
            # 转换为字符串（15位左右）
            return str(id_value)
    
    def generate_id_with_prefix(self, prefix: str = "") -> str:
        """
        生成带前缀的ID
        
        Args:
            prefix: ID前缀
            
        Returns:
            str: 带前缀的ID
        """
        id_value = self.generate_id()
        return f"{prefix}{id_value}" if prefix else id_value

# 全局雪花算法生成器实例
_snowflake_generator: Optional[SnowflakeGenerator] = None
_generator_lock = threading.Lock()

def get_snowflake_generator(datacenter_id: int = 1, worker_id: int = 1) -> SnowflakeGenerator:
    """
    获取雪花算法生成器实例（单例模式）
    
    Args:
        datacenter_id: 数据中心ID
        worker_id: 工作机器ID
        
    Returns:
        SnowflakeGenerator: 雪花算法生成器实例
    """
    global _snowflake_generator
    
    if _snowflake_generator is None:
        with _generator_lock:
            if _snowflake_generator is None:
                _snowflake_generator = SnowflakeGenerator(datacenter_id, worker_id)
    
    return _snowflake_generator

def generate_id() -> str:
    """
    生成雪花算法ID
    
    Returns:
        str: 15位左右的唯一ID
    """
    generator = get_snowflake_generator()
    return generator.generate_id()

def generate_id_with_prefix(prefix: str = "") -> str:
    """
    生成带前缀的雪花算法ID
    
    Args:
        prefix: ID前缀
        
    Returns:
        str: 带前缀的ID
    """
    generator = get_snowflake_generator()
    return generator.generate_id_with_prefix(prefix) 